package com.bruce.asurada.controller;

import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.bruce.asurada.assistant.Assistant;
import com.bruce.asurada.assistant.StreamAssistant;
import com.bruce.asurada.bean.ChatForm;
import com.bruce.asurada.common.Result;
import com.bruce.asurada.dto.ChatMessageDTO;
import com.bruce.asurada.store.MongoChatMemotyStore;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.Content;
import dev.langchain4j.data.message.TextContent;
import jakarta.annotation.Resource;
import jakarta.validation.constraints.NotBlank;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;


@RestController
@Slf4j
@RequestMapping("/assistant")
public class AssistantController {


    @Resource
    private  Assistant assistant;

    @Resource
    private MongoChatMemotyStore mongoChatMemotyStore;


    @Resource
    private StreamAssistant streamAssistant;


    /**
     * 获取聊天记录
     * @param memoryId
     * @return
     */
    @GetMapping("/getChatHistory")
    public Result<List<ChatMessageDTO>>  getChatHistory(@RequestParam @NotBlank(message = "聊天记录ID不能为空") String memoryId) {
        List<ChatMessage> messages = mongoChatMemotyStore.getMessages(memoryId);
        log.info("memoryId={}, messages={}", memoryId, messages);
        if(messages == null || messages.size() <= 0){
            return Result.error("没有聊天记录");
        }

        List<ChatMessageDTO> chatMessageDTOS = new ArrayList<>() ;
        for (ChatMessage message : messages) {
            String sender = null ;
            String content = null ;
            if(message instanceof UserMessage){
                  sender = "USER";
                  List<Content> contents = ((UserMessage) message).contents() ;
                  content = generateUserContent(contents) ;
            }
            if(message instanceof AiMessage){
                  sender = "AI";
                  content = ((AiMessage) message).text() ;
            }
            if(StringUtils.isEmpty(sender) || StringUtils.isEmpty(content)){
                continue;
            }
            ChatMessageDTO chatMessageDTO = ChatMessageDTO.builder()
                    .sender(sender)
                    .content(content)
                    .build();
            chatMessageDTOS.add(chatMessageDTO);
        }
        return Result.success(chatMessageDTOS) ;
    }
    

    /**
     * 聊天
     * @param chatForm
     * @return
     */
    @PostMapping(value ="/chat")
    public String chat(@RequestBody ChatForm chatForm) {
        log.info("chat memoryId :{}",chatForm.getMemoryId());
        return assistant.chat(chatForm.getMemoryId(),chatForm.getMessages());
    }


    /**
     * 流式输出聊天
     * @param chatForm
     * @return
     */
    @PostMapping(value ="/streamChat",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamChat(@RequestBody ChatForm chatForm) {
        return streamAssistant.chat(chatForm.getMemoryId(),chatForm.getMessages());
    }

    private String generateUserContent(List<Content> contents){
        if(contents == null || contents.isEmpty()){
            return null ;
        }

        StringBuilder contentStr = new StringBuilder();
        for (Content content : contents) {
            if(content instanceof TextContent){
                TextContent con = (TextContent) content;
                contentStr.append(con.text());
            }
        }
        return contentStr.toString() ;
    }

}
